# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: daniel

import collections
import os
import matplotlib.pyplot as plt
import pandas as pd
from absl import app, flags


FLAGS = flags.FLAGS


class HttpRequestTimeoutAnalyzer(object):
  def __init__(self, og_log_config):
    self._og_log_config = og_log_config
    # self._http_request_timeout_data = collections.namedtuple(
    #     'HttpRequestTimeoutData',
    #     ['machine', 'strategy_name', 'trading_date', "details", "summary"]
    # )
    # self._http_request_timeout_data.machine = og_log_config.machine
    # self._http_request_timeout_data.strategy_name = og_log_config.strategy_name
    # self._http_request_timeout_data.trading_date = og_log_config.trading_date
    # self._http_request_timeout_data.details = {}
    # self._http_request_timeout_data.summary = []
    self._http_request_timeout_data = collections.defaultdict(dict)
    self._http_request_timeout_data["machine"] = og_log_config.machine
    self._http_request_timeout_data["strategy_name"] = og_log_config.strategy_name
    self._http_request_timeout_data["trading_date"] = og_log_config.trading_date
    self._http_request_timeout_data["summary"] = []
    self._http_request_timeout_data["details"] = {}

  def filter_telemetry_order_latency_records(self, records):
    ret = []
    for record in records:
      if record['type'] != 'TELEMETRY':
        continue
      if 'latency' not in record['telemetry']:
        continue
      if record['telemetry']['type'] != 'LATENCY' or \
         record['telemetry']['latency']['tag'] != 'ORDER_LATENCY':
        continue
      ret.append(record)
    return ret

  def retrieve_http_request_timeout_record(self,
                                           records,
                                           start_dt,
                                           end_dt):
    tmp = []
    for record in records:
      mark_points = record['telemetry']['latency']['mark_points']
      is_os_rest_timeout_error = False
      timestamp = 0
      for mark_point in mark_points:
        if mark_point['point_name'] == 'OG_REST_TIMEOUT_ERROR':
          is_os_rest_timeout_error = True
          timestamp = mark_point['timestamp']
          break
      if is_os_rest_timeout_error is False:
        continue
      rec = {
          'timestamp': int(timestamp),
          'num_of_to': 1,
      }
      tmp.append(rec)
    num = len(tmp)
    ret = []
    if num == 0:
      return num, ret
    # generate # of timeout per 5 mins
    x_beg = int(pd.to_datetime(start_dt.strftime('%Y%m%dT%H%M%S')).timestamp()) * 1000000000
    x_end = int(pd.to_datetime(end_dt.strftime('%Y%m%dT%H%M%S')).timestamp()) * 1000000000
    x = x_beg
    y = 0
    while x_beg < x_end:
      for each in tmp:
        if each['timestamp'] < x_beg:
          continue
        if each['timestamp'] >= x_beg + 300000000000:
          continue
        y = y + 1
      info = {'timestamp': x, 'num_of_to': y}
      ret.append(info)
      x_beg = x_beg + 300000000000
      x = x_beg
      y = 0

    return num, sorted(ret, key=lambda x: x['timestamp'])


  def process(self, records):
    start_dt = self._og_log_config.start_dt
    end_dt = self._og_log_config.end_dt
    strategy_name = self._og_log_config.strategy_name
    telemetry_order_latency_records = self.filter_telemetry_order_latency_records(records)
    # ignore strategy with zero http conn timeout
    if len(telemetry_order_latency_records) == 0:
      return self._http_request_timeout_data

    http_request_timeout_total_num, http_request_timeout_records = \
        self.retrieve_http_request_timeout_record(telemetry_order_latency_records, start_dt, end_dt)
    http_request_timeout_details = {
        'http_request_timeout_total_num': http_request_timeout_total_num,
        'http_request_timeout_records': http_request_timeout_records,
    }
    self._http_request_timeout_data['details'] = http_request_timeout_details

    http_request_timeout_summary = []
    info = {
        'strategy_name': strategy_name,
        'total_num': http_request_timeout_total_num,
    }
    http_request_timeout_summary.append(info)
    self._http_request_timeout_data['summary'] = http_request_timeout_summary

    return self._http_request_timeout_data


class HttpRequestTimeoutPlotter(object):
  def __init__(self, plotter_config):
    self._plotter_config = plotter_config

  # def store_strategy_analysis_result(self):
  #   tmp_dir = f"{self._tmp_dir}/{self._trading_date}"
  #   if not os.path.exists(tmp_dir):
  #     os.makedirs(tmp_dir)
  #   combine_report = {
  #       "rate_limit_hit_summary": self._rate_limit_hit_summary,
  #       "http_request_timeout_summary": self._http_request_timeout_summary,
  #       "fail_orders_summary": self._fail_orders_summary,
  #       "error_stats_summary": self._error_stats_summary,
  #   }
  #   report = f"{tmp_dir}/{self._strategy_name}_{self._machine}.json"
  #   with open(report, 'w') as ofile:
  #     json.dump(combine_report, ofile, indent=2)

  def generate_http_request_timeout_summary_report(self,
                                                   axe,
                                                   summary,
                                                   trade_date,
                                                   output_dir):
    if len(summary) == 0:
      return

    records = sorted(summary, key=lambda x: x['total_num'], reverse=True)
    label = []
    total = []
    for record in records:
      label.append(record['strategy_name'])
      total.append(record['total_num'])

    # if plt_in is None:
    #   plt = plt_local
    # else:
    #   plt = plt_in
    # plt.rcParams.update({'font.size': 7})
    #fig = plt.figure(figsize=(30, 16), dpi=200)
    # draw bar chart
    #plt.subplot()
    axe.bar(label, total, width=0.35, label='')
    #axe.ylabel('total # of timedout')
    #axe.xticks(range(len(label)), label, rotation=270)

    #axe.title('http request timeout summary')
    #axe.grid(True)
    #axe.legend()
    #fig.tight_layout()
    #output_path = f"{output_dir}/{trade_date}/http_request_timeout"
    #if not os.path.exists(output_path):
    #  os.makedirs(output_path)
    #output_file = f"{trade_date}_http_request_timeout_summary_report.svg"
    #print(f"write to {output_path}/{output_file}")
    #fig.savefig(f"{output_path}/{output_file}", dpi=fig.dpi)
    #if plt_in is None:
    #  plt.close(fig)

  def generate_http_request_timeout_report(self,
                                           details,
                                           strategy_name,
                                           machine,
                                           start_dt,
                                           end_dt,
                                           output_dir):
    if not details:
      return

    plt.rcParams.update({'font.size': 7})
    fig = plt.figure(figsize=(10, 6), dpi=200)
    # draw linear chart
    plt.subplot()
    df = pd.DataFrame(details['http_request_timeout_records'])
    df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
    plt.plot(df.timestamp, df.num_of_to, marker='.', markersize=0.1, label='')
    plt.title('# of http request timeout in 5 mins')

    plt.grid(True)
    plt.legend()
    fig.tight_layout(pad=1.0, w_pad=1.0, h_pad=2.0)
    trade_date = start_dt.strftime('%Y%m%d')
    start_dt_str = start_dt.strftime('%Y%m%dT%H%M%S')
    end_dt_str = end_dt.strftime('%Y%m%dT%H%M%S')
    output_path = f"{output_dir}/{trade_date}/http_request_timeout/{strategy_name}/{machine}"
    if not os.path.exists(output_path):
      os.makedirs(output_path)
    output_file = f"{start_dt_str}-{end_dt_str}.svg"
    print(f"write to {output_path}/{output_file}")
    fig.savefig(f"{output_path}/{output_file}", dpi=fig.dpi)
    plt.close(fig)

  def plot_summary(self, result, axe):
    if result is None:
      return
    output_dir = self._plotter_config.output_dir
    trade_date = self._plotter_config.trading_date
    self.generate_http_request_timeout_summary_report(axe,
                                                      result['summary'],
                                                      trade_date,
                                                      output_dir)

  def plot_details(self, result):
    if result is None:
      return
    output_dir = self._plotter_config.output_dir
    trade_date = self._plotter_config.trading_date
    machine = self._plotter_config.machine
    strategy_name = self._plotter_config.strategy_name
    start_dt = self._plotter_config.start_dt
    end_dt = self._plotter_config.end_dt
    self.generate_http_request_timeout_report(result['details'],
                                              strategy_name,
                                              machine,
                                              start_dt,
                                              end_dt,
                                              output_dir)


def main(_):
  # TODO(daniel):
  print("@daniel: to impliment")


if __name__ == '__main__':
  flags.DEFINE_string('root_dir',
                      '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log',
                      'root dir specify')
  flags.DEFINE_string('output_dir',
                      '/remote/iosg/strat-1/buckets/analysis.derived.coin/live/plot/og_log_plot',
                      'output_dir.')
  flags.DEFINE_string('tmp_dir',
                      '/tmp',
                      'tmp dir for slurm run')
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('machine', None, 'Instance machine name.')
  flags.DEFINE_string('strategy_name', None, 'Strategy name.')
  flags.DEFINE_string('start_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('end_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('symbol', None, '')
  flags.DEFINE_string('log_type', None, '')
  flags.DEFINE_bool('local_run', False, "")
  flags.DEFINE_bool('slurm_run', False, "")
  flags.DEFINE_bool('generate_summary_report', False, "generate summary report only")
  app.run(main)
